package com.sc.base.common.interceptor.sql;

import com.alibaba.druid.pool.DruidPooledPreparedStatement;
import com.alibaba.druid.pool.PreparedStatementHolder;
import com.alibaba.druid.proxy.jdbc.ConnectionProxy;
import com.alibaba.druid.proxy.jdbc.PreparedStatementProxyImpl;
import com.google.common.collect.LinkedHashMultimap;
import com.sc.base.common.log.SQLRespLog;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.connection.MasterSlaveConnection;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.statement.MasterSlavePreparedStatement;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.util.Collection;
import java.util.Map;


/**
 * 获取执行的sql信息并打印所属数据源
 * @author liuwei
 * @time 2021-1-14 17:46:22
 * */
public class SQLPerInterceptorUtils {

    public static void handleDataSourceKey(SQLRespLog sqlRespLog, PreparedStatement preparedStatement){
        // TODO DruidPooled
        if(preparedStatement instanceof DruidPooledPreparedStatement){
            DruidPooledPreparedStatement druidPooledPreparedStatement = (DruidPooledPreparedStatement)preparedStatement;
            PreparedStatementHolder preparedStatementHolder = druidPooledPreparedStatement.getPreparedStatementHolder();
            PreparedStatement statementObj = preparedStatementHolder.getStatement();
            if(statementObj instanceof PreparedStatementProxyImpl){
                PreparedStatementProxyImpl statement = (PreparedStatementProxyImpl)preparedStatementHolder.getStatement();
                ConnectionProxy connectionProxy = statement.getConnectionProxy();
                String dataSourceKey = connectionProxy.getDirectDataSource().getName();
                String dbType = connectionProxy.getDirectDataSource().getDbType();
                sqlRespLog.setDataSourceKey(dataSourceKey);
                sqlRespLog.setDatabaseType(dbType);
            }
        }
        // TODO ShardingJdbc
        if(preparedStatement instanceof MasterSlavePreparedStatement){
            MasterSlavePreparedStatement masterSlavePreparedStatement = (MasterSlavePreparedStatement)preparedStatement;
            MasterSlaveConnection connection = masterSlavePreparedStatement.getConnection();
            LinkedHashMultimap<String, Connection> cachedConnectionsMap = (LinkedHashMultimap)connection.getCachedConnections();
            Map<String, Collection<Connection>> stringCollectionMap = cachedConnectionsMap.asMap();
            String dataSourceKey = stringCollectionMap.size()>0 ? stringCollectionMap.keySet().stream().findFirst().get():"";
            sqlRespLog.setDataSourceKey(dataSourceKey);
        }
    }
}
